// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/function_vector.h>

namespace kumo::pollux::functions {
    // Returns the total number of nested elements for the specified top-levels rows
    // in array or map vector. T is either ArrayVector or MapVector.
    template<typename T>
    vector_size_t countElements(
        const SelectivityVector &rows,
        DecodedVector &decodedVector) {
        auto indices = decodedVector.indices();
        auto rawSizes = decodedVector.base()->as<T>()->rawSizes();

        vector_size_t count = 0;
        rows.applyToSelected([&](vector_size_t row) {
            if (decodedVector.is_null_at(row)) {
                return;
            }
            // In some cases, the array/map vector is wrapped in dictionary encoding
            // that can explode the cardinality of the underlying elements. This check
            // helps ensure we fail instead of silently wrapping around and causing
            // memory corruption or segfaults.
            count = checkedPlus<vector_size_t>(count, rawSizes[indices[row]]);
        });
        return count;
    }

    // Returns an array of indices that allows aligning captures with the nested
    // elements of an array or vector. For each top-level row, the index equal to
    // the row number is repeated for each of the nested rows.
    template<typename T>
    BufferPtr toWrapCapture(
        vector_size_t size,
        const Callable *callable,
        const SelectivityVector &topLevelRows,
        const std::shared_ptr<T> &topLevelVector) {
        if (!callable->hasCapture()) {
            return nullptr;
        }

        auto rawNulls = topLevelVector->raw_nulls();
        auto rawSizes = topLevelVector->rawSizes();
        auto rawOffsets = topLevelVector->rawOffsets();

        BufferPtr wrapCapture = allocate_indices(size, topLevelVector->pool());
        auto rawWrapCapture = wrapCapture->asMutable<vector_size_t>();
        topLevelRows.applyToSelected([&](vector_size_t row) {
            if (rawNulls && bits::isBitNull(rawNulls, row)) {
                return;
            }
            auto size = rawSizes[row];
            auto offset = rawOffsets[row];
            for (auto i = 0; i < size; ++i) {
                rawWrapCapture[offset + i] = row;
            }
        });
        return wrapCapture;
    }

    // Given possibly wrapped array vector, flattens the wrappings and returns a
    // flat array vector. Returns the original vector unmodified if the vector is
    // not wrapped. Flattening is shallow, e.g. elements vector may still be
    // wrapped.
    ArrayVectorPtr flattenArray(
        const SelectivityVector &rows,
        const VectorPtr &vector,
        DecodedVector &decodedVector);

    // Given possibly wrapped map vector, flattens the wrappings and returns a flat
    // map vector. Returns the original vector unmodified if the vector is not
    // wrapped. Flattening is shallow, e.g. keys and values vectors may still be
    // wrapped.
    MapVectorPtr flattenMap(
        const SelectivityVector &rows,
        const VectorPtr &vector,
        DecodedVector &decodedVector);

    // Creates a nulls buffer to hold rows.size() nulls. Copies the original nulls
    // in 'vector' for positions [rows.begin(), rows.end()) and sets nulls for
    // unselected rows in 'rows'.
    BufferPtr addNullsForUnselectedRows(
        const VectorPtr &vector,
        const SelectivityVector &rows);
} // namespace kumo::pollux::functions
